Problem

FusedRMSNorm (and the sibling fused layer/RMS norm custom ops) leak two CUDA tensors per forward call under torch.no_grad(), as reported in #1999. On the torch.library.custom_op path (PyTorch >= 2.4) the setup_context functions unconditionally call save_for_backward; those saved tensors are retained in autograd metadata that is not released after a no_grad forward, leaking the saved activation + invvar each call.

Fix

In each affected setup_context (apex/normalization/fused_layer_norm.py), assign the scalar ctx fields first, then return early when torch.is_grad_enabled() is False, skipping save_for_backward. Backward can never run under no_grad, so nothing is lost, and the grad-enabled training path is unchanged.

Testing / verification status

Runtime-unverified — this was developed on a machine without a CUDA GPU, so the leak reproducer was not executed. The change compiles (py_compile) and passes ruff, and the reasoning is that skipping the save under no_grad removes the retained references regardless of the exact internal mechanism. I'd appreciate a maintainer with a GPU running the issue's count_cuda_tensors reproducer to confirm delta == 0 before merge; happy to adjust if the leak originates elsewhere.

Developed with AI assistance.

Addresses #1999